import logging
from typing import List, Tuple, Union, Optional, Any

from app.utils.cv_tools import show_image_with_roi
from app.utils.file_tools import extract_filename_without_extension, build_output_path_with_page_info, clear_directory
from app.utils.pymupdf_tools import pdf_get_content_region, crop_pdf_with_white_mask, get_pages_from_range
from app.utils.timing import timing_decorator
from tests.base_test import base_test_case

logger = base_test_case.get_logger(__name__)
# import logging
# print(f"输出日志Name：{logging.Logger.manager.loggerDict.keys()}")
TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir

"""公式识别"""


def simple_test(input_file):
    from paddleocr import FormulaRecognition
    model = FormulaRecognition(model_name="PP-FormulaNet_plus-M")
    result = model.predict(input=input_file, batch_size=1)
    for res in result:
        res.print()
        res.save_to_img(save_path="./output/formula/")
        res.save_to_json(save_path="./output/formula/")
    pass


# 输入pdf文件和 page_range 范围
def pdf_with_page_test(source, zoom_factor=1.0, page_range: Optional[Union[Tuple[int, int], List[int]]] = None):
    # 如果输入类型为str或文件对象，提取文件名，不包括后缀，并返回文件名，否则返回None
    file_name = extract_filename_without_extension(source)
    images = get_pages_from_range(source, page_range=page_range)
    # 查看原图
    for image in images:
        show_image_with_roi(image, zoom_factor=zoom_factor, is_show=True)

    result = []
    # from paddleocr import FormulaRecognition
    # model = FormulaRecognition(model_name="PP-FormulaNet_plus-M")
    # result = model.predict(input=images, batch_size=1)

    # 获取实际处理的页面索引列表
    if page_range is not None:
        if isinstance(page_range, tuple):
            # page_range 是一个范围 (start, end)
            processed_pages = list(range(page_range[0], page_range[1]))
        else:
            # page_range 是一个列表 [page1, page2, ...]
            processed_pages = page_range
    else:
        # 没有指定 page_range，页面索引从 0 开始
        processed_pages = list(range(len(images)))

    for i, res in enumerate(result):
        res.print()
        # 使用新函数构建输出路径
        output_path = build_output_path_with_page_info(source, processed_pages[i], "./output/formula")
        res.save_to_img(save_path=output_path)
        res.save_to_json(save_path=output_path)
    pass


def get_pdf_with_page_layout(source, zoom_factor=1.0,
                             page_range: Optional[Union[Tuple[int, int], List[int]]] = None,
                             is_show_log=False, roi=None, output_prefix="test",
                             is_crop_pdf_with_white_mask=False,  # 是否生成裁切后的pdf
                             ):
    # # 获取索引页
    # page_ranges = get_pages_from_range(source, page_range=page_range)
    # print(page_ranges)
    # # 返回正文roi
    # rois = pdf_get_content_region_v2(source, zoom_factor=zoom_factor, page_range=page_range, is_show_log=is_show_log)
    # print(rois)
    # for page, roi in zip(page_ranges, rois):
    #     show_image_with_roi(page, zoom_factor=zoom_factor, is_show=False, roi=roi)

    # 按索引页重构 pdf
    output_pdf_path = OUTPUT_DATA_DIR / f"{output_prefix}_crop.pdf"

    # 定义检测器函数（带参数）
    def my_roi_detector(page):
        return pdf_get_content_region(page, zoom_factor=zoom_factor, is_show_log=False)

    # 裁切 pdf
    crop_pdf_with_white_mask(input_path=source, page_range=page_range, roi_detector=my_roi_detector, output_path=output_pdf_path)

    # 输入pdf文件,已经剪切过,去除注释内容的pdf文件
    images = get_pages_from_range(str(output_pdf_path), page_range=page_range)
    instance = create_layout_detection_instance()
    output_dir = build_output_path_with_page_info(output_pdf_path)
    # 先删除输出目录下的所有文件
    clear_directory(output_dir)
    for page_idx, image in enumerate(images):
        results = layout_detection_predict(image, instance, params)
        for i, res in enumerate(results):
            # # 添加按cls_id升序排序的代码
            # # res.get("boxes", []).sort(key=lambda x: x['cls_id'])
            #
            # # 按照从上到下、从左到右的顺序排序坐标
            # # 创建副本避免修改原始数据
            # boxes_copy = res.get("boxes", []).copy()
            # boxes_copy.sort(key=lambda x: (x['coordinate'][1], x['coordinate'][0]))
            # # 使用排序后的数据进行后续处理
            # res_sorted = res.copy()
            # res_sorted["boxes"] = boxes_copy

            # 添加按cls_id升序排序的代码
            res.get("boxes", []).sort(key=lambda x: x['cls_id'])
            # 按照从上到下、从左到右的顺序排序坐标
            res.get("boxes", []).sort(key=lambda x: (x['coordinate'][1], x['coordinate'][0]))
            res["page_index"] = page_idx

            # res.print()
            res.save_to_img(save_path=f"{output_dir}/{page_idx}.jpg")
            res.save_to_json(save_path=f"{output_dir}/{page_idx}.json")
            # # 保存按坐标排序的结果
            # import json
            # from pathlib import Path
            # import numpy as np
            #
            # # 确保输出目录存在
            # Path(output_dir).mkdir(parents=True, exist_ok=True)
            # # 移除
            # res_sorted.pop("input_img", None)
            #
            # # 递归转换函数，将numpy数组和特殊数值类型转换为可JSON序列化的格式
            # def convert(o):
            #     if isinstance(o, np.ndarray):
            #         return o.tolist()
            #     if isinstance(o, np.integer):
            #         return int(o)
            #     if isinstance(o, np.floating):
            #         return float(o)
            #     if isinstance(o, dict):
            #         return {k: convert(v) for k, v in o.items()}
            #     if isinstance(o, (list, tuple)):
            #         return [convert(x) for x in o]
            #     return o
            #
            # # 保存排序后的结果为JSON文件
            # with open(f"{output_dir}/{page_idx}_sorted.json", 'w', encoding='utf-8') as f:
            #     json.dump({"res": convert(res_sorted)}, f, indent=4, ensure_ascii=False)
            # 提取公式坐标


# 初始化
init_kwargs = {
    # 'img_size': [545, 755], # 初始化不支持 img_size
    'device': 'cpu', 'cpu_threads': 10,
    # "threshold": 0.5,
    # "model_dir": "", # /home/tomcat/.paddlex/official_models/PP-DocLayout_plus-L
    # "model_name": "layout",
    # 'enable_hpi': False,
    # 'enable_mkldnn': True,
    # 'mkldnn_cache_capacity': 10,
    # 'precision': 'fp32',
    # 'use_tensorrt': False
}


# 创建实例
@timing_decorator(prefix="创建实例", log_level=logging.INFO)
def create_layout_detection_instance():
    from paddleocr import LayoutDetection
    instance = LayoutDetection(**init_kwargs)
    logger.info(f"{instance.__class__.__name__} instance created")
    return instance


params = {
    "batch_size": 10,
    # "threshold": 0.5,
    # 'layout_nms': 'bool|None',
    # 'layout_unclip_ratio': 'float|list|dict|None',
    # 'layout_merge_bboxes_mode': 'str|dict|None',
    # 预测时,不支持img_size=[545,755]
}


@timing_decorator(prefix="OCR处理", log_level=logging.INFO, error_log_level=logging.WARNING)
def layout_detection_predict(images: str | list, instance, params: dict[str, int | float | str]) -> list[Any]:
    results = instance.predict(images, **params)
    return results


if __name__ == '__main__':
    # build_output_path_with_page_info(str(TEST_DATA_DIR / "1711605374231.pdf"))
    page_range = (0, 1)
    # simple_test(str(TEST_DATA_DIR / "formula.png"))
    # # 测试查看 str(TEST_DATA_DIR / "1711605374231.pdf") 5-6正文区域与page_range范围，返回正文roi
    # rois = pdf_get_content_region_v2(str(TEST_DATA_DIR / "1711605374231.pdf"), page_range=page_range, is_show_log=False)
    # print(rois)
    #
    # # 测试输入str文件与page_range范围，返回page范围
    # page_ranges = get_pages_from_range(str(TEST_DATA_DIR / "1711605374231.pdf"), page_range=page_range)
    # print(page_ranges)
    # for page, roi in zip(page_ranges, rois):
    #     show_image_with_roi(page, zoom_factor=1.0, is_show=True, roi=roi)

    # 先做版面检查,提取公式坐标
    # get_pdf_with_page_layout(str(TEST_DATA_DIR / "1711605374231.pdf"), page_range=page_range, is_show_log=False, output_prefix="1711605374231_02")
    # get_pdf_with_page_layout(str(TEST_DATA_DIR / "1715339805571.pdf"), page_range=(5,6), is_show_log=False, output_prefix="1715339805571")
    # get_pdf_with_page_layout(str(TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf"), page_range=None, is_show_log=False,
    #                          output_prefix="25-注会-轻1-财务成本管理[上册](第3章)")

    # 分页处理
    get_pdf_with_page_layout(str(TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf"),
                             page_range=page_range, is_show_log=False,
                             output_prefix="0-1",
                             is_crop_pdf_with_white_mask=False,
                             )

    # # 识别 1715339805571.pdf文件 4-5页面的公式
    # pdf_with_page_test(str(TEST_DATA_DIR / "1711605374231.pdf"), page_range=page_range)
    pass
